import math
import numpy as np
import torch
import torch.nn as nn
from torch.nn import init
from torch.nn import functional as F

from utils import log_normal_density


class CNNNetwork(nn.Module):
    def __init__(self, frames, action_space):
        super(CNNNetwork, self).__init__()
        self.logstd = nn.Parameter(torch.zeros(action_space))

        self.conv1 = nn.Conv2d(frames, 16, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1)
        self.conv3 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        self.relu = nn.ReLU()
        self.act_fc1 = nn.Linear(64*60*80, 256)
        self.act_fc2 =  nn.Linear(256+7+6, 128)
        self.actor1 = nn.Linear(128, 3)
        self.actor2 = nn.Linear(128, 1)

        self.conv1_critic = nn.Conv2d(1, 16, kernel_size=3, padding=1)
        self.conv2_critic = nn.Conv2d(16, 32, kernel_size=3, padding=1)
        self.conv3_critic = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.crt_fc1 = nn.Linear(64*60*80, 256)
        self.crt_fc2 = nn.Linear(256+7+6, 128)
        self.critic = nn.Linear(128, 1)



    def forward(self, x, goal, speed):
        """
            returns value estimation, action, log_action_prob
        """
        if len(x.shape)==3 and len(goal.shape)==1 and len(speed.shape)==1:
            x = x.unsqueeze(0)
            goal = goal.unsqueeze(0)
            speed = speed.unsqueeze(0)
        # action
        a = self.relu(self.conv1(x))
        a = self.pool(a)
        a = self.relu(self.conv2(a))
        a = self.pool(a)
        a = self.relu(self.conv3(a))
        a = self.pool(a)

        a = a.view(a.shape[0], -1)
        a = F.relu(self.act_fc1(a))

        a = torch.cat((a, goal, speed), dim=-1)
        a = F.relu(self.act_fc2(a))
        # TODO: scaling action?
        mean1 = F.tanh(self.actor1(a))
        mean2 = F.tanh(self.actor2(a))
        mean = torch.cat((mean1, mean2), dim=-1)

        logstd = self.logstd.expand_as(mean)
        std = torch.exp(logstd)
        action = torch.normal(mean, std)

        # action prob on log scale
        logprob = log_normal_density(action, mean, std=std, log_std=logstd)

        # value
        v = self.relu(self.conv1(x))
        v = self.pool(v)
        v = self.relu(self.conv2(v))
        v = self.pool(v)
        v = self.relu(self.conv3(v))
        v = self.pool(v)
        v = v.view(v.shape[0], -1)
        v = F.relu(self.crt_fc1(v))
        v = torch.cat((v, goal, speed), dim=-1)
        v = F.relu(self.crt_fc2(v))
        v = self.critic(v)


        return v, action, logprob, mean

    def evaluate_actions(self, x, goal, speed, action):
        v, _, _, mean = self.forward(x, goal, speed)
        logstd = self.logstd.expand_as(mean)
        std = torch.exp(logstd)
        # evaluate
        logprob = log_normal_density(action, mean, log_std=logstd, std=std)
        dist_entropy = 0.5 + 0.5 * math.log(2 * math.pi) + logstd
        dist_entropy = dist_entropy.sum(-1).mean()
        return v, logprob, dist_entropy


if __name__ == '__main__':
    from torch.autograd import Variable

    import time
    start_time = time.time()

    net = CNNNetwork(frames=1, action_space=4)

    observation = Variable(torch.randn(32, 1, 480, 640))
    goal = torch.randn(32, 7)
    vel = torch.randn(32, 6)
    v, action, logprob, mean = net.forward(observation,goal,vel)

    end_time = time.time()
    print(end_time-start_time)

